

<!DOCTYPE html>
<html class="writer-html5" lang="en" >
<head>
  <meta charset="utf-8" />
  
  <meta name="viewport" content="width=device-width, initial-scale=1.0" />
  
  <title>mindspore.nn.layer.math &mdash; MindSpore master documentation</title>
  

  
  <link rel="stylesheet" href="../../../../_static/css/theme.css" type="text/css" />
  <link rel="stylesheet" href="../../../../_static/pygments.css" type="text/css" />

  
  

  
  

  

  
  <!--[if lt IE 9]>
    <script src="../../../../_static/js/html5shiv.min.js"></script>
  <![endif]-->
  
    
      <script type="text/javascript" id="documentation_options" data-url_root="../../../../" src="../../../../_static/documentation_options.js"></script>
        <script src="../../../../_static/jquery.js"></script>
        <script src="../../../../_static/underscore.js"></script>
        <script src="../../../../_static/doctools.js"></script>
        <script src="../../../../_static/language_data.js"></script>
        <script async="async" src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/latest.js?config=TeX-AMS-MML_HTMLorMML"></script>
    
    <script type="text/javascript" src="../../../../_static/js/theme.js"></script>

    
    <link rel="index" title="Index" href="../../../../genindex.html" />
    <link rel="search" title="Search" href="../../../../search.html" /> 
</head>

<body class="wy-body-for-nav">

   
  <div class="wy-grid-for-nav">
    
    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
      <div class="wy-side-scroll">
        <div class="wy-side-nav-search" >
          

          
            <a href="../../../../index.html" class="icon icon-home"> MindSpore
          

          
          </a>

          
            
            
          

          
<div role="search">
  <form id="rtd-search-form" class="wy-form" action="../../../../search.html" method="get">
    <input type="text" name="q" placeholder="Search docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>

          
        </div>

        
        <div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
          
            
            
              
            
            
              <p class="caption"><span class="caption-text">MindSpore Python API</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.html">mindspore</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.common.initializer.html">mindspore.common.initializer</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.communication.html">mindspore.communication</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.compression.html">mindspore.compression</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.context.html">mindspore.context</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.html">mindspore.dataset</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.audio.html">mindspore.dataset.audio</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.config.html">mindspore.dataset.config</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.text.html">mindspore.dataset.text</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.transforms.html">mindspore.dataset.transforms</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.dataset.vision.html">mindspore.dataset.vision</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.mindrecord.html">mindspore.mindrecord</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.html">mindspore.nn</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.probability.html">mindspore.nn.probability</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.nn.transformer.html">mindspore.nn.transformer</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.numpy.html">mindspore.numpy</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.ops.html">mindspore.ops</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.parallel.html">mindspore.parallel</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.parallel.nn.html">mindspore.parallel.nn</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.profiler.html">mindspore.profiler</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.scipy.html">mindspore.scipy</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.train.html">mindspore.train</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../../api_python/mindspore.boost.html">mindspore.boost</a></li>
</ul>
<p class="caption"><span class="caption-text">MindSpore C++ API</span></p>
<ul>
<li class="toctree-l1"><a class="reference external" href="https://www.mindspore.cn/lite/api/zh-CN/master/api_cpp/mindspore.html">MindSpore Lite↗</a></li>
</ul>

            
          
        </div>
        
      </div>
    </nav>

    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap">

      
      <nav class="wy-nav-top" aria-label="top navigation">
        
          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
          <a href="../../../../index.html">MindSpore</a>
        
      </nav>


      <div class="wy-nav-content">
        
        <div class="rst-content">
        
          

















<div role="navigation" aria-label="breadcrumbs navigation">

  <ul class="wy-breadcrumbs">
    
      <li><a href="../../../../index.html" class="icon icon-home"></a> &raquo;</li>
        
          <li><a href="../../../index.html">Module code</a> &raquo;</li>
        
      <li>mindspore.nn.layer.math</li>
    
    
      <li class="wy-breadcrumbs-aside">
        
      </li>
    
  </ul>

  
  <hr/>
</div>
          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
           <div itemprop="articleBody">
            
  <h1>Source code for mindspore.nn.layer.math</h1><div class="highlight"><pre>
<span></span><span class="c1"># Copyright 2020-2021 Huawei Technologies Co., Ltd</span>
<span class="c1">#</span>
<span class="c1"># Licensed under the Apache License, Version 2.0 (the &quot;License&quot;);</span>
<span class="c1"># you may not use this file except in compliance with the License.</span>
<span class="c1"># You may obtain a copy of the License at</span>
<span class="c1">#</span>
<span class="c1"># http://www.apache.org/licenses/LICENSE-2.0</span>
<span class="c1">#</span>
<span class="c1"># Unless required by applicable law or agreed to in writing, software</span>
<span class="c1"># distributed under the License is distributed on an &quot;AS IS&quot; BASIS,</span>
<span class="c1"># WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.</span>
<span class="c1"># See the License for the specific language governing permissions and</span>
<span class="c1"># limitations under the License.</span>
<span class="c1"># ============================================================================</span>
<span class="sd">&quot;&quot;&quot;math&quot;&quot;&quot;</span>
<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
<span class="kn">from</span> <span class="nn">mindspore.ops</span> <span class="kn">import</span> <span class="n">operations</span> <span class="k">as</span> <span class="n">P</span>
<span class="kn">from</span> <span class="nn">mindspore.common.tensor</span> <span class="kn">import</span> <span class="n">Tensor</span>
<span class="kn">from</span> <span class="nn">mindspore.common._decorator</span> <span class="kn">import</span> <span class="n">deprecated</span>
<span class="kn">from</span> <span class="nn">mindspore.ops.primitive</span> <span class="kn">import</span> <span class="n">constexpr</span>
<span class="kn">from</span> <span class="nn">mindspore.ops</span> <span class="kn">import</span> <span class="n">functional</span> <span class="k">as</span> <span class="n">F</span>
<span class="kn">from</span> <span class="nn">..cell</span> <span class="kn">import</span> <span class="n">Cell</span>
<span class="kn">from</span> <span class="nn">...common</span> <span class="kn">import</span> <span class="n">dtype</span> <span class="k">as</span> <span class="n">mstype</span>
<span class="kn">from</span> <span class="nn">..._checkparam</span> <span class="kn">import</span> <span class="n">Validator</span> <span class="k">as</span> <span class="n">validator</span>

<span class="n">__all__</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;ReduceLogSumExp&#39;</span><span class="p">,</span>
           <span class="s1">&#39;Range&#39;</span><span class="p">,</span>
           <span class="s1">&#39;LGamma&#39;</span><span class="p">,</span>
           <span class="s1">&#39;DiGamma&#39;</span><span class="p">,</span>
           <span class="s1">&#39;IGamma&#39;</span><span class="p">,</span>
           <span class="s1">&#39;LBeta&#39;</span><span class="p">,</span>
           <span class="s1">&#39;MatMul&#39;</span><span class="p">,</span>
           <span class="s1">&#39;Moments&#39;</span><span class="p">,</span>
           <span class="s1">&#39;MatInverse&#39;</span><span class="p">,</span>
           <span class="s1">&#39;MatDet&#39;</span><span class="p">,</span>
           <span class="p">]</span>

<span class="n">_BASE_LANCZOS_COEFF</span> <span class="o">=</span> <span class="mf">0.99999999999980993227684700473478</span>
<span class="n">_LANCZOS_COEFFICIENTS</span> <span class="o">=</span> <span class="p">[</span><span class="mf">676.520368121885098567009190444019</span><span class="p">,</span>
                         <span class="o">-</span><span class="mf">1259.13921672240287047156078755283</span><span class="p">,</span>
                         <span class="mf">771.3234287776530788486528258894</span><span class="p">,</span>
                         <span class="o">-</span><span class="mf">176.61502916214059906584551354</span><span class="p">,</span>
                         <span class="mf">12.507343278686904814458936853</span><span class="p">,</span>
                         <span class="o">-</span><span class="mf">0.13857109526572011689554707</span><span class="p">,</span>
                         <span class="mf">9.984369578019570859563e-6</span><span class="p">,</span>
                         <span class="mf">1.50563273514931155834e-7</span><span class="p">]</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">_check_input_dtype</span><span class="p">(</span><span class="n">param_name</span><span class="p">,</span> <span class="n">input_dtype</span><span class="p">,</span> <span class="n">allow_dtypes</span><span class="p">,</span> <span class="n">cls_name</span><span class="p">):</span>
    <span class="n">validator</span><span class="o">.</span><span class="n">check_type_name</span><span class="p">(</span><span class="n">param_name</span><span class="p">,</span> <span class="n">input_dtype</span><span class="p">,</span> <span class="n">allow_dtypes</span><span class="p">,</span> <span class="n">cls_name</span><span class="p">)</span>


<div class="viewcode-block" id="ReduceLogSumExp"><a class="viewcode-back" href="../../../../api_python/nn/mindspore.nn.ReduceLogSumExp.html#mindspore.nn.ReduceLogSumExp">[docs]</a><span class="k">class</span> <span class="nc">ReduceLogSumExp</span><span class="p">(</span><span class="n">Cell</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Reduces a dimension of a tensor by calculating exponential for all elements in the dimension,</span>
<span class="sd">    then calculate logarithm of the sum.</span>

<span class="sd">    .. math::</span>

<span class="sd">        ReduceLogSumExp(x) = \log(\sum(e^x))</span>

<span class="sd">    Args:</span>
<span class="sd">        axis (Union[int, tuple(int), list(int)]) - The dimensions to reduce. Default: (), reduce all dimensions.</span>
<span class="sd">            Only constant value is allowed.</span>
<span class="sd">        keep_dims (bool): If True, keep these reduced dimensions and the length is 1.</span>
<span class="sd">            If False, don&#39;t keep these dimensions.</span>
<span class="sd">            Default : False.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **x** (Tensor) - The input tensor. With float16 or float32 data type.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Tensor, has the same dtype as the `x`.</span>

<span class="sd">        - If axis is (), and keep_dims is False,</span>
<span class="sd">          the output is a 0-D tensor representing the sum of all elements in the input tensor.</span>
<span class="sd">        - If axis is int, set as 2, and keep_dims is False,</span>
<span class="sd">          the shape of output is :math:`(x_1, x_3, ..., x_R)`.</span>
<span class="sd">        - If axis is tuple(int), set as (2, 3), and keep_dims is False,</span>
<span class="sd">          the shape of output is :math:`(x_1, x_4, ..., x_R)`.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If `axis` is not one of int, list, tuple.</span>
<span class="sd">        TypeError: If `keep_dims` is not bool.</span>
<span class="sd">        TypeError: If dtype of `x` is neither float16 nor float32.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU`` ``CPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; x = Tensor(np.random.randn(3, 4, 5, 6).astype(np.float32))</span>
<span class="sd">        &gt;&gt;&gt; op = nn.ReduceLogSumExp(1, keep_dims=True)</span>
<span class="sd">        &gt;&gt;&gt; output = op(x)</span>
<span class="sd">        &gt;&gt;&gt; print(output.shape)</span>
<span class="sd">        (3, 1, 5, 6)</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">axis</span><span class="p">,</span> <span class="n">keep_dims</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize ReduceLogSumExp.&quot;&quot;&quot;</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">ReduceLogSumExp</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;axis&#39;</span><span class="p">,</span> <span class="n">axis</span><span class="p">,</span> <span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">list</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;keep_dims&#39;</span><span class="p">,</span> <span class="n">keep_dims</span><span class="p">,</span> <span class="p">[</span><span class="nb">bool</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">axis</span> <span class="o">=</span> <span class="n">axis</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">exp</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Exp</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">sum</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">ReduceSum</span><span class="p">(</span><span class="n">keep_dims</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">log</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Log</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">construct</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span>
        <span class="n">exp</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
        <span class="n">sumexp</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">exp</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">axis</span><span class="p">)</span>
        <span class="n">logsumexp</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">sumexp</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">logsumexp</span></div>


<div class="viewcode-block" id="Range"><a class="viewcode-back" href="../../../../api_python/nn/mindspore.nn.Range.html#mindspore.nn.Range">[docs]</a><span class="k">class</span> <span class="nc">Range</span><span class="p">(</span><span class="n">Cell</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Creates a sequence of numbers in range [start, limit) with step size delta.</span>

<span class="sd">    The size of output is :math:`\left \lfloor \frac{limit-start}{delta}  \right \rfloor + 1` and `delta` is the gap</span>
<span class="sd">    between two values in the tensor.</span>

<span class="sd">    .. math::</span>

<span class="sd">        out_{i+1} = out_{i} +delta</span>

<span class="sd">    Args:</span>
<span class="sd">        start (Union[int, float]): If `limit` is `None`, the value acts as limit in the range and first entry</span>
<span class="sd">            defaults to `0`. Otherwise, it acts as first entry in the range.</span>
<span class="sd">        limit (Union[int, float]): Acts as upper limit of sequence. If `None`, defaults to the value of `start`</span>
<span class="sd">            while set the first entry of the range to `0`. It can not be equal to `start`. Default: None.</span>
<span class="sd">        delta (Union[int, float]): Increment of the range. It can not be equal to zero. Default: 1.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Tensor, the dtype is int if the dtype of `start`, `limit` and `delta` all are int. Otherwise, dtype is float.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU`` ``CPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; net = nn.Range(1, 8, 2)</span>
<span class="sd">        &gt;&gt;&gt; output = net()</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [1 3 5 7]</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">start</span><span class="p">,</span> <span class="n">limit</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">delta</span><span class="o">=</span><span class="mi">1</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize Range.&quot;&quot;&quot;</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">Range</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
        <span class="k">if</span> <span class="n">delta</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="si">}</span><span class="s2">&#39;, the &#39;delta&#39; can not be zero.&quot;</span><span class="p">)</span>
        <span class="n">data</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">start</span><span class="p">,</span> <span class="n">limit</span><span class="p">,</span> <span class="n">delta</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">data</span><span class="o">.</span><span class="n">dtype</span> <span class="o">==</span> <span class="n">np</span><span class="o">.</span><span class="n">float</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">ms_dtype</span> <span class="o">=</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float32</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="bp">self</span><span class="o">.</span><span class="n">ms_dtype</span> <span class="o">=</span> <span class="n">mstype</span><span class="o">.</span><span class="n">int32</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">result_tensor</span> <span class="o">=</span> <span class="n">Tensor</span><span class="p">(</span><span class="n">data</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">ms_dtype</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">construct</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">result_tensor</span></div>


<span class="k">class</span> <span class="nc">LGamma</span><span class="p">(</span><span class="n">Cell</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Calculates LGamma using Lanczos&#39; approximation referring to &quot;A Precision Approximation of the Gamma Function&quot;.</span>
<span class="sd">    The algorithm is:</span>

<span class="sd">    .. math::</span>
<span class="sd">        \begin{array}{ll} \\</span>
<span class="sd">            lgamma(z + 1) = \frac{(\log(2) + \log(pi))}{2} + (z + 1/2) * log(t(z)) - t(z) + A(z) \\</span>
<span class="sd">            t(z) = z + kLanczosGamma + 1/2 \\</span>
<span class="sd">            A(z) = kBaseLanczosCoeff + \sum_{k=1}^n \frac{kLanczosCoefficients[i]}{z + k}</span>
<span class="sd">        \end{array}</span>

<span class="sd">    However, if the input is less than 0.5 use Euler&#39;s reflection formula:</span>

<span class="sd">    .. math::</span>

<span class="sd">        lgamma(x) = \log(pi) - lgamma(1-x) - \log(abs(sin(pi * x)))</span>

<span class="sd">    And please note that</span>

<span class="sd">    .. math::</span>

<span class="sd">        lgamma(+/-inf) = +inf</span>

<span class="sd">    Thus, the behaviour of LGamma follows:</span>

<span class="sd">    - when x &gt; 0.5, return log(Gamma(x))</span>
<span class="sd">    - when x &lt; 0.5 and is not an integer, return the real part of Log(Gamma(x)) where Log is the complex logarithm</span>
<span class="sd">    - when x is an integer less or equal to 0, return +inf</span>
<span class="sd">    - when x = +/- inf, return +inf</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **x** (Tensor) - The input tensor. Only float16, float32 are supported.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Tensor, has the same shape and dtype as the `x`.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If dtype of `x` is neither float16 nor float32.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; x = Tensor(np.array([2, 3, 4]).astype(np.float32))</span>
<span class="sd">        &gt;&gt;&gt; op = nn.LGamma()</span>
<span class="sd">        &gt;&gt;&gt; output = op(x)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [3.5762787e-07 6.9314754e-01 1.7917603e+00]</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize LGamma.&quot;&quot;&quot;</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">LGamma</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
        <span class="c1"># const numbers</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">k_lanczos_gamma</span> <span class="o">=</span> <span class="mi">7</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">k_base_lanczos_coeff</span> <span class="o">=</span> <span class="n">_BASE_LANCZOS_COEFF</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">k_lanczos_coefficients</span> <span class="o">=</span> <span class="n">_LANCZOS_COEFFICIENTS</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">one_half</span> <span class="o">=</span> <span class="mf">0.5</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">one</span> <span class="o">=</span> <span class="mi">1</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">two</span> <span class="o">=</span> <span class="mi">2</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">inf</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">inf</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">pi</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">pi</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">log_2</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">two</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">log_pi</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">pi</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">log_sqrt_two_pi</span> <span class="o">=</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">log_2</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">log_pi</span><span class="p">)</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">two</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lanczos_gamma_plus_one_half</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">k_lanczos_gamma</span> <span class="o">+</span> <span class="mf">0.5</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">log_lanczos_gamma_plus_one_half</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">lanczos_gamma_plus_one_half</span><span class="p">)</span>

        <span class="c1"># operations</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">log</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Log</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">log1p</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Log1p</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">abs</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Abs</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">shape</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Shape</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">DType</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">fill</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Fill</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">floor</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Floor</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">equal</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Equal</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">greater</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Greater</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">less</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Less</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lessequal</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">LessEqual</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">select</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Select</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">sin</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Sin</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">isfinite</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">IsFinite</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">construct</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span>
        <span class="n">input_dtype</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
        <span class="n">_check_input_dtype</span><span class="p">(</span><span class="s2">&quot;x&quot;</span><span class="p">,</span> <span class="n">input_dtype</span><span class="p">,</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">float16</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="n">infinity</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">fill</span><span class="p">(</span><span class="n">input_dtype</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="bp">self</span><span class="o">.</span><span class="n">inf</span><span class="p">)</span>

        <span class="n">need_to_reflect</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">less</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="mf">0.5</span><span class="p">)</span>
        <span class="n">neg_input</span> <span class="o">=</span> <span class="o">-</span><span class="n">x</span>
        <span class="n">z</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">select</span><span class="p">(</span><span class="n">need_to_reflect</span><span class="p">,</span> <span class="n">neg_input</span><span class="p">,</span> <span class="n">x</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span>

        <span class="nd">@constexpr</span>
        <span class="k">def</span> <span class="nf">_calculate_reflected_x</span><span class="p">(</span><span class="n">z</span><span class="p">,</span> <span class="n">k_base_lanczos_coeff</span><span class="p">,</span> <span class="n">k_lanczos_coefficients</span><span class="p">):</span>
            <span class="n">reflex_x</span> <span class="o">=</span> <span class="n">k_base_lanczos_coeff</span>
            <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">8</span><span class="p">):</span>
                <span class="n">product_</span> <span class="o">=</span> <span class="n">k_lanczos_coefficients</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">/</span> <span class="p">(</span><span class="n">z</span> <span class="o">+</span> <span class="n">i</span> <span class="o">+</span> <span class="mi">1</span><span class="p">)</span>
                <span class="n">reflex_x</span> <span class="o">=</span> <span class="n">product_</span> <span class="o">+</span> <span class="n">reflex_x</span>
            <span class="k">return</span> <span class="n">reflex_x</span>
        <span class="n">reflex_x</span> <span class="o">=</span> <span class="n">_calculate_reflected_x</span><span class="p">(</span><span class="n">z</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">k_base_lanczos_coeff</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">k_lanczos_coefficients</span><span class="p">)</span>

        <span class="n">t</span> <span class="o">=</span> <span class="n">z</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">lanczos_gamma_plus_one_half</span>
        <span class="n">log_t</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">log1p</span><span class="p">(</span><span class="n">z</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">lanczos_gamma_plus_one_half</span><span class="p">)</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">log_lanczos_gamma_plus_one_half</span>

        <span class="n">log_y</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">reflex_x</span><span class="p">)</span> <span class="o">+</span> <span class="p">(</span><span class="n">z</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">one_half</span> <span class="o">-</span> <span class="n">t</span> <span class="o">/</span> <span class="n">log_t</span><span class="p">)</span> <span class="o">*</span> <span class="n">log_t</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">log_sqrt_two_pi</span>

        <span class="n">abs_input</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">abs</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
        <span class="n">abs_frac_input</span> <span class="o">=</span> <span class="n">abs_input</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">floor</span><span class="p">(</span><span class="n">abs_input</span><span class="p">)</span>
        <span class="n">x</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">select</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">lessequal</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="mf">0.0</span><span class="p">),</span> <span class="bp">self</span><span class="o">.</span><span class="n">select</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">equal</span><span class="p">(</span><span class="n">abs_frac_input</span><span class="p">,</span> <span class="mf">0.0</span><span class="p">),</span> <span class="n">infinity</span><span class="p">,</span> <span class="n">x</span><span class="p">),</span> <span class="n">x</span><span class="p">)</span>
        <span class="n">reduced_frac_input</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">select</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">greater</span><span class="p">(</span><span class="n">abs_frac_input</span><span class="p">,</span> <span class="mf">0.5</span><span class="p">),</span>
                                         <span class="mi">1</span> <span class="o">-</span> <span class="n">abs_frac_input</span><span class="p">,</span> <span class="n">abs_frac_input</span><span class="p">)</span>
        <span class="n">reflection_denom</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">sin</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">pi</span> <span class="o">*</span> <span class="n">reduced_frac_input</span><span class="p">))</span>

        <span class="n">reflection</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">select</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">isfinite</span><span class="p">(</span><span class="n">reflection_denom</span><span class="p">),</span>
                                 <span class="o">-</span><span class="n">reflection_denom</span> <span class="o">-</span> <span class="n">log_y</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">log_pi</span><span class="p">,</span> <span class="c1"># pylint: disable=invalid-unary-operand-type</span>
                                 <span class="o">-</span><span class="n">reflection_denom</span><span class="p">)</span> <span class="c1"># pylint: disable=invalid-unary-operand-type</span>

        <span class="n">result</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">select</span><span class="p">(</span><span class="n">need_to_reflect</span><span class="p">,</span> <span class="n">reflection</span><span class="p">,</span> <span class="n">log_y</span><span class="p">)</span>

        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">select</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">isfinite</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="n">result</span><span class="p">,</span> <span class="n">infinity</span><span class="p">)</span>


<span class="k">class</span> <span class="nc">DiGamma</span><span class="p">(</span><span class="n">Cell</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Calculates Digamma using Lanczos&#39; approximation referring to &quot;A Precision Approximation of the Gamma Function&quot;.</span>
<span class="sd">    The algorithm is:</span>

<span class="sd">    .. math::</span>
<span class="sd">        \begin{array}{ll} \\</span>
<span class="sd">            digamma(z + 1) = log(t(z)) + A&#39;(z) / A(z) - kLanczosGamma / t(z) \\</span>
<span class="sd">            t(z) = z + kLanczosGamma + 1/2 \\</span>
<span class="sd">            A(z) = kBaseLanczosCoeff + \sum_{k=1}^n \frac{kLanczosCoefficients[i]}{z + k} \\</span>
<span class="sd">            A&#39;(z) = \sum_{k=1}^n \frac{kLanczosCoefficients[i]}{{z + k}^2}</span>
<span class="sd">        \end{array}</span>

<span class="sd">    However, if the input is less than 0.5 use Euler&#39;s reflection formula:</span>

<span class="sd">    .. math::</span>

<span class="sd">        digamma(x) = digamma(1 - x) - pi * cot(pi * x)</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **x** (Tensor[Number]) - The input tensor. Only float16, float32 are supported.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Tensor, has the same shape and dtype as the `x`.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If dtype of `x` is neither float16 nor float32.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; x = Tensor(np.array([2, 3, 4]).astype(np.float32))</span>
<span class="sd">        &gt;&gt;&gt; op = nn.DiGamma()</span>
<span class="sd">        &gt;&gt;&gt; output = op(x)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [0.42278463  0.92278427 1.2561178]</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize DiGamma.&quot;&quot;&quot;</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">DiGamma</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
        <span class="c1"># const numbers</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">k_lanczos_gamma</span> <span class="o">=</span> <span class="mi">7</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">k_base_lanczos_coeff</span> <span class="o">=</span> <span class="n">_BASE_LANCZOS_COEFF</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">k_lanczos_coefficients</span> <span class="o">=</span> <span class="n">_LANCZOS_COEFFICIENTS</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">nan</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">nan</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">pi</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">pi</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lanczos_gamma_plus_one_half</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">k_lanczos_gamma</span> <span class="o">+</span> <span class="mf">0.5</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">log_lanczos_gamma_plus_one_half</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">lanczos_gamma_plus_one_half</span><span class="p">)</span>

        <span class="c1"># operations</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">log1p</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Log1p</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">abs</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Abs</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">shape</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Shape</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">DType</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">fill</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Fill</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">floor</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Floor</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">equal</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Equal</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">less</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Less</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">select</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Select</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">sin</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Sin</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">cos</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Cos</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">logicaland</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">LogicalAnd</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">construct</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span>
        <span class="n">input_dtype</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
        <span class="n">_check_input_dtype</span><span class="p">(</span><span class="s2">&quot;x&quot;</span><span class="p">,</span> <span class="n">input_dtype</span><span class="p">,</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">float16</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="n">need_to_reflect</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">less</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="mf">0.5</span><span class="p">)</span>
        <span class="n">neg_input</span> <span class="o">=</span> <span class="o">-</span><span class="n">x</span>
        <span class="n">z</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">select</span><span class="p">(</span><span class="n">need_to_reflect</span><span class="p">,</span> <span class="n">neg_input</span><span class="p">,</span> <span class="n">x</span> <span class="o">-</span> <span class="mi">1</span><span class="p">)</span>

        <span class="nd">@constexpr</span>
        <span class="k">def</span> <span class="nf">_calculate_num_denom</span><span class="p">(</span><span class="n">z</span><span class="p">,</span> <span class="n">k_base_lanczos_coeff</span><span class="p">,</span> <span class="n">k_lanczos_coefficients</span><span class="p">):</span>
            <span class="n">num</span> <span class="o">=</span> <span class="mi">0</span>
            <span class="n">denom</span> <span class="o">=</span> <span class="n">k_base_lanczos_coeff</span>
            <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">8</span><span class="p">):</span>
                <span class="n">num</span> <span class="o">=</span> <span class="n">num</span> <span class="o">-</span> <span class="n">k_lanczos_coefficients</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">/</span> <span class="p">((</span><span class="n">z</span> <span class="o">+</span> <span class="n">i</span> <span class="o">+</span> <span class="mi">1</span><span class="p">)</span> <span class="o">*</span> <span class="p">(</span><span class="n">z</span> <span class="o">+</span> <span class="n">i</span> <span class="o">+</span> <span class="mi">1</span><span class="p">))</span>
                <span class="n">denom</span> <span class="o">=</span> <span class="n">denom</span> <span class="o">+</span> <span class="n">k_lanczos_coefficients</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">/</span> <span class="p">(</span><span class="n">z</span> <span class="o">+</span> <span class="n">i</span> <span class="o">+</span> <span class="mi">1</span><span class="p">)</span>
            <span class="k">return</span> <span class="n">num</span><span class="p">,</span> <span class="n">denom</span>
        <span class="n">num</span><span class="p">,</span> <span class="n">denom</span> <span class="o">=</span> <span class="n">_calculate_num_denom</span><span class="p">(</span><span class="n">z</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">k_base_lanczos_coeff</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">k_lanczos_coefficients</span><span class="p">)</span>

        <span class="n">t</span> <span class="o">=</span> <span class="n">z</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">lanczos_gamma_plus_one_half</span>
        <span class="n">log_t</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">log1p</span><span class="p">(</span><span class="n">z</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">lanczos_gamma_plus_one_half</span><span class="p">)</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">log_lanczos_gamma_plus_one_half</span>

        <span class="n">y</span> <span class="o">=</span> <span class="n">log_t</span> <span class="o">+</span> <span class="n">num</span> <span class="o">/</span> <span class="n">denom</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">k_lanczos_gamma</span> <span class="o">/</span> <span class="n">t</span>

        <span class="n">reduced_input</span> <span class="o">=</span> <span class="n">x</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">abs</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">floor</span><span class="p">(</span><span class="n">x</span> <span class="o">+</span> <span class="mf">0.5</span><span class="p">))</span>
        <span class="n">reflection</span> <span class="o">=</span> <span class="n">y</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">pi</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">cos</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">pi</span> <span class="o">*</span> <span class="n">reduced_input</span><span class="p">)</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">sin</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">pi</span> <span class="o">*</span> <span class="n">reduced_input</span><span class="p">)</span>
        <span class="n">real_result</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">select</span><span class="p">(</span><span class="n">need_to_reflect</span><span class="p">,</span> <span class="n">reflection</span><span class="p">,</span> <span class="n">y</span><span class="p">)</span>
        <span class="n">nan</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">fill</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="n">np</span><span class="o">.</span><span class="n">nan</span><span class="p">)</span>

        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">select</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">logicaland</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">less</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="mi">0</span><span class="p">),</span> <span class="bp">self</span><span class="o">.</span><span class="n">equal</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">floor</span><span class="p">(</span><span class="n">x</span><span class="p">))),</span>
                           <span class="n">nan</span><span class="p">,</span> <span class="n">real_result</span><span class="p">)</span>


<span class="n">eps_fp32</span> <span class="o">=</span> <span class="n">Tensor</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">finfo</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span><span class="o">.</span><span class="n">eps</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">_while_helper_func</span><span class="p">(</span><span class="n">cond</span><span class="p">,</span> <span class="n">body</span><span class="p">,</span> <span class="n">vals</span><span class="p">):</span>
    <span class="k">while</span> <span class="n">cond</span><span class="p">(</span><span class="n">vals</span><span class="p">)</span><span class="o">.</span><span class="n">any</span><span class="p">():</span>
        <span class="n">vals</span> <span class="o">=</span> <span class="n">body</span><span class="p">(</span><span class="n">vals</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">vals</span>


<span class="k">def</span> <span class="nf">_igamma_series</span><span class="p">(</span><span class="n">ax</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">a</span><span class="p">,</span> <span class="n">enabled</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Helper function for computing Igamma using a power series.&quot;&quot;&quot;</span>

    <span class="n">logicaland</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">LogicalAnd</span><span class="p">()</span>
    <span class="n">greater</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Greater</span><span class="p">()</span>
    <span class="n">fill</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Fill</span><span class="p">()</span>
    <span class="n">shape</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Shape</span><span class="p">()</span>
    <span class="n">dtype</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">DType</span><span class="p">()</span>
    <span class="n">select</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Select</span><span class="p">()</span>

    <span class="c1"># If more data types are supported, this epsilon need to be selected.</span>
    <span class="n">epsilon</span> <span class="o">=</span> <span class="n">eps_fp32</span>

    <span class="k">def</span> <span class="nf">cond</span><span class="p">(</span><span class="n">vals</span><span class="p">):</span>
        <span class="n">enabled</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
        <span class="k">return</span> <span class="n">enabled</span>

    <span class="k">def</span> <span class="nf">body</span><span class="p">(</span><span class="n">vals</span><span class="p">):</span>
        <span class="n">enabled</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
        <span class="n">r</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
        <span class="n">c</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span>
        <span class="n">ans</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">3</span><span class="p">]</span>
        <span class="n">x</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">4</span><span class="p">]</span>
        <span class="n">dc_da</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">5</span><span class="p">]</span>
        <span class="n">dans_da</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">6</span><span class="p">]</span>

        <span class="n">r</span> <span class="o">=</span> <span class="n">r</span> <span class="o">+</span> <span class="mi">1</span>
        <span class="n">dc_da</span> <span class="o">=</span> <span class="n">dc_da</span> <span class="o">*</span> <span class="p">(</span><span class="n">x</span> <span class="o">/</span> <span class="n">r</span><span class="p">)</span> <span class="o">+</span> <span class="p">(</span><span class="o">-</span><span class="mi">1</span> <span class="o">*</span> <span class="n">c</span> <span class="o">*</span> <span class="n">x</span><span class="p">)</span> <span class="o">/</span> <span class="p">(</span><span class="n">r</span> <span class="o">*</span> <span class="n">r</span><span class="p">)</span>
        <span class="n">dans_da</span> <span class="o">=</span> <span class="n">dans_da</span> <span class="o">+</span> <span class="n">dc_da</span>
        <span class="n">c</span> <span class="o">=</span> <span class="n">c</span> <span class="o">*</span> <span class="p">(</span><span class="n">x</span> <span class="o">/</span> <span class="n">r</span><span class="p">)</span>
        <span class="n">ans</span> <span class="o">=</span> <span class="n">ans</span> <span class="o">+</span> <span class="n">c</span>
        <span class="n">conditional</span> <span class="o">=</span> <span class="n">logicaland</span><span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">greater</span><span class="p">(</span><span class="n">c</span> <span class="o">/</span> <span class="n">ans</span><span class="p">,</span> <span class="n">epsilon</span><span class="p">))</span>

        <span class="k">return</span> <span class="p">(</span><span class="n">conditional</span><span class="p">,</span> <span class="n">select</span><span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">r</span><span class="p">,</span> <span class="n">vals</span><span class="p">[</span><span class="mi">1</span><span class="p">]),</span>
                <span class="n">select</span><span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">c</span><span class="p">,</span> <span class="n">vals</span><span class="p">[</span><span class="mi">2</span><span class="p">]),</span> <span class="n">select</span><span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">ans</span><span class="p">,</span> <span class="n">vals</span><span class="p">[</span><span class="mi">3</span><span class="p">]),</span>
                <span class="n">select</span><span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">vals</span><span class="p">[</span><span class="mi">4</span><span class="p">]),</span> <span class="n">select</span><span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">dc_da</span><span class="p">,</span> <span class="n">vals</span><span class="p">[</span><span class="mi">5</span><span class="p">]),</span>
                <span class="n">select</span><span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">dans_da</span><span class="p">,</span> <span class="n">vals</span><span class="p">[</span><span class="mi">6</span><span class="p">]))</span>

    <span class="n">ones</span> <span class="o">=</span> <span class="n">fill</span><span class="p">(</span><span class="n">dtype</span><span class="p">(</span><span class="n">a</span><span class="p">),</span> <span class="n">shape</span><span class="p">(</span><span class="n">a</span><span class="p">),</span> <span class="mi">1</span><span class="p">)</span>
    <span class="n">zeros</span> <span class="o">=</span> <span class="n">fill</span><span class="p">(</span><span class="n">dtype</span><span class="p">(</span><span class="n">a</span><span class="p">),</span> <span class="n">shape</span><span class="p">(</span><span class="n">a</span><span class="p">),</span> <span class="mi">0</span><span class="p">)</span>
    <span class="n">vals</span> <span class="o">=</span> <span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">a</span><span class="p">,</span> <span class="n">ones</span><span class="p">,</span> <span class="n">ones</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">zeros</span><span class="p">,</span> <span class="n">zeros</span><span class="p">)</span>

    <span class="n">vals</span> <span class="o">=</span> <span class="n">_while_helper_func</span><span class="p">(</span><span class="n">cond</span><span class="p">,</span> <span class="n">body</span><span class="p">,</span> <span class="n">vals</span><span class="p">)</span>
    <span class="n">ans</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">3</span><span class="p">]</span>
    <span class="k">return</span> <span class="p">(</span><span class="n">ans</span> <span class="o">*</span> <span class="n">ax</span><span class="p">)</span> <span class="o">/</span> <span class="n">a</span>


<span class="k">def</span> <span class="nf">_igammac_continued_fraction</span><span class="p">(</span><span class="n">ax</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">a</span><span class="p">,</span> <span class="n">enabled</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;Helper function for computing Igammac using a continued fraction.&quot;&quot;&quot;</span>

    <span class="n">abs_x</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Abs</span><span class="p">()</span>
    <span class="n">logicaland</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">LogicalAnd</span><span class="p">()</span>
    <span class="n">greater</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Greater</span><span class="p">()</span>
    <span class="n">less</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Less</span><span class="p">()</span>
    <span class="n">notequal</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">NotEqual</span><span class="p">()</span>
    <span class="n">fill</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Fill</span><span class="p">()</span>
    <span class="n">shape</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Shape</span><span class="p">()</span>
    <span class="n">dtype</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">DType</span><span class="p">()</span>
    <span class="n">select</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Select</span><span class="p">()</span>

    <span class="c1"># If more data types are supported, this epsilon need to be selected.</span>
    <span class="n">epsilon</span> <span class="o">=</span> <span class="n">eps_fp32</span>

    <span class="k">def</span> <span class="nf">cond</span><span class="p">(</span><span class="n">vals</span><span class="p">):</span>
        <span class="n">enabled</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
        <span class="n">c</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">5</span><span class="p">]</span>
        <span class="k">return</span> <span class="n">logicaland</span><span class="p">(</span><span class="n">less</span><span class="p">(</span><span class="n">c</span><span class="p">,</span> <span class="mi">2000</span><span class="p">),</span> <span class="n">enabled</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">body</span><span class="p">(</span><span class="n">vals</span><span class="p">):</span>
        <span class="n">enabled</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
        <span class="n">ans</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
        <span class="n">t</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">2</span><span class="p">]</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">3</span><span class="p">]</span>
        <span class="n">z</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">4</span><span class="p">]</span>
        <span class="n">c</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">5</span><span class="p">]</span>
        <span class="n">pkm1</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">6</span><span class="p">]</span>
        <span class="n">qkm1</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">7</span><span class="p">]</span>
        <span class="n">pkm2</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">8</span><span class="p">]</span>
        <span class="n">qkm2</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">9</span><span class="p">]</span>

        <span class="n">dpkm2_da</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">10</span><span class="p">]</span>
        <span class="n">dqkm2_da</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">11</span><span class="p">]</span>
        <span class="n">dpkm1_da</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">12</span><span class="p">]</span>
        <span class="n">dqkm1_da</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">13</span><span class="p">]</span>
        <span class="n">dans_da</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">14</span><span class="p">]</span>

        <span class="n">c</span> <span class="o">=</span> <span class="n">c</span> <span class="o">+</span> <span class="mi">1</span>
        <span class="n">y</span> <span class="o">=</span> <span class="n">y</span> <span class="o">+</span> <span class="mi">1</span>
        <span class="n">z</span> <span class="o">=</span> <span class="n">z</span> <span class="o">+</span> <span class="mi">2</span>

        <span class="n">yc</span> <span class="o">=</span> <span class="n">y</span> <span class="o">*</span> <span class="n">c</span>
        <span class="n">pk</span> <span class="o">=</span> <span class="n">pkm1</span> <span class="o">*</span> <span class="n">z</span> <span class="o">-</span> <span class="n">pkm2</span> <span class="o">*</span> <span class="n">yc</span>
        <span class="n">qk</span> <span class="o">=</span> <span class="n">qkm1</span> <span class="o">*</span> <span class="n">z</span> <span class="o">-</span> <span class="n">qkm2</span> <span class="o">*</span> <span class="n">yc</span>
        <span class="n">qk_is_nonzero</span> <span class="o">=</span> <span class="n">notequal</span><span class="p">(</span><span class="n">qk</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
        <span class="n">r</span> <span class="o">=</span> <span class="n">pk</span> <span class="o">/</span> <span class="n">qk</span>

        <span class="n">t</span> <span class="o">=</span> <span class="n">select</span><span class="p">(</span><span class="n">qk_is_nonzero</span><span class="p">,</span> <span class="n">abs_x</span><span class="p">((</span><span class="n">ans</span> <span class="o">-</span> <span class="n">r</span><span class="p">)</span> <span class="o">/</span> <span class="n">r</span><span class="p">),</span> <span class="n">fill</span><span class="p">(</span><span class="n">dtype</span><span class="p">(</span><span class="n">t</span><span class="p">),</span> <span class="n">shape</span><span class="p">(</span><span class="n">t</span><span class="p">),</span> <span class="mi">1</span><span class="p">))</span>
        <span class="n">ans</span> <span class="o">=</span> <span class="n">select</span><span class="p">(</span><span class="n">qk_is_nonzero</span><span class="p">,</span> <span class="n">r</span><span class="p">,</span> <span class="n">ans</span><span class="p">)</span>

        <span class="n">dpk_da</span> <span class="o">=</span> <span class="n">dpkm1_da</span> <span class="o">*</span> <span class="n">z</span> <span class="o">-</span> <span class="n">pkm1</span> <span class="o">-</span> <span class="n">dpkm2_da</span> <span class="o">*</span> <span class="n">yc</span> <span class="o">+</span> <span class="n">pkm2</span> <span class="o">*</span> <span class="n">c</span>
        <span class="n">dqk_da</span> <span class="o">=</span> <span class="n">dqkm1_da</span> <span class="o">*</span> <span class="n">z</span> <span class="o">-</span> <span class="n">qkm1</span> <span class="o">-</span> <span class="n">dqkm2_da</span> <span class="o">*</span> <span class="n">yc</span> <span class="o">+</span> <span class="n">qkm2</span> <span class="o">*</span> <span class="n">c</span>
        <span class="n">dans_da_new</span> <span class="o">=</span> <span class="n">select</span><span class="p">(</span><span class="n">qk_is_nonzero</span><span class="p">,</span> <span class="p">(</span><span class="n">dpk_da</span> <span class="o">-</span> <span class="n">ans</span> <span class="o">*</span> <span class="n">dqk_da</span><span class="p">)</span> <span class="o">/</span> <span class="n">qk</span><span class="p">,</span> <span class="n">dans_da</span><span class="p">)</span>
        <span class="n">grad_conditional</span> <span class="o">=</span> <span class="n">select</span><span class="p">(</span><span class="n">qk_is_nonzero</span><span class="p">,</span>
                                  <span class="n">abs_x</span><span class="p">(</span><span class="n">dans_da_new</span> <span class="o">-</span> <span class="n">dans_da</span><span class="p">),</span>
                                  <span class="n">fill</span><span class="p">(</span><span class="n">dtype</span><span class="p">(</span><span class="n">dans_da</span><span class="p">),</span> <span class="n">shape</span><span class="p">(</span><span class="n">dans_da</span><span class="p">),</span> <span class="mi">1</span><span class="p">))</span>

        <span class="n">pkm2</span> <span class="o">=</span> <span class="n">pkm1</span>
        <span class="n">pkm1</span> <span class="o">=</span> <span class="n">pk</span>
        <span class="n">qkm2</span> <span class="o">=</span> <span class="n">qkm1</span>
        <span class="n">qkm1</span> <span class="o">=</span> <span class="n">qk</span>

        <span class="n">dpkm2_da</span> <span class="o">=</span> <span class="n">dpkm1_da</span>
        <span class="n">dqkm2_da</span> <span class="o">=</span> <span class="n">dqkm1_da</span>
        <span class="n">dpkm1_da</span> <span class="o">=</span> <span class="n">dpk_da</span>
        <span class="n">dqkm1_da</span> <span class="o">=</span> <span class="n">dqk_da</span>

        <span class="n">rescale</span> <span class="o">=</span> <span class="n">greater</span><span class="p">(</span><span class="n">abs_x</span><span class="p">(</span><span class="n">pk</span><span class="p">),</span> <span class="mi">1</span> <span class="o">/</span> <span class="n">epsilon</span><span class="p">)</span>
        <span class="n">pkm2</span> <span class="o">=</span> <span class="n">select</span><span class="p">(</span><span class="n">rescale</span><span class="p">,</span> <span class="n">pkm2</span> <span class="o">*</span> <span class="n">epsilon</span><span class="p">,</span> <span class="n">pkm2</span><span class="p">)</span>
        <span class="n">pkm1</span> <span class="o">=</span> <span class="n">select</span><span class="p">(</span><span class="n">rescale</span><span class="p">,</span> <span class="n">pkm1</span> <span class="o">*</span> <span class="n">epsilon</span><span class="p">,</span> <span class="n">pkm1</span><span class="p">)</span>
        <span class="n">qkm2</span> <span class="o">=</span> <span class="n">select</span><span class="p">(</span><span class="n">rescale</span><span class="p">,</span> <span class="n">qkm2</span> <span class="o">*</span> <span class="n">epsilon</span><span class="p">,</span> <span class="n">qkm2</span><span class="p">)</span>
        <span class="n">qkm1</span> <span class="o">=</span> <span class="n">select</span><span class="p">(</span><span class="n">rescale</span><span class="p">,</span> <span class="n">qkm1</span> <span class="o">*</span> <span class="n">epsilon</span><span class="p">,</span> <span class="n">qkm1</span><span class="p">)</span>

        <span class="n">dpkm2_da</span> <span class="o">=</span> <span class="n">select</span><span class="p">(</span><span class="n">rescale</span><span class="p">,</span> <span class="n">dpkm2_da</span> <span class="o">*</span> <span class="n">epsilon</span><span class="p">,</span> <span class="n">dpkm2_da</span><span class="p">)</span>
        <span class="n">dqkm2_da</span> <span class="o">=</span> <span class="n">select</span><span class="p">(</span><span class="n">rescale</span><span class="p">,</span> <span class="n">dqkm2_da</span> <span class="o">*</span> <span class="n">epsilon</span><span class="p">,</span> <span class="n">dqkm2_da</span><span class="p">)</span>
        <span class="n">dpkm1_da</span> <span class="o">=</span> <span class="n">select</span><span class="p">(</span><span class="n">rescale</span><span class="p">,</span> <span class="n">dpkm1_da</span> <span class="o">*</span> <span class="n">epsilon</span><span class="p">,</span> <span class="n">dpkm1_da</span><span class="p">)</span>
        <span class="n">dqkm1_da</span> <span class="o">=</span> <span class="n">select</span><span class="p">(</span><span class="n">rescale</span><span class="p">,</span> <span class="n">dqkm1_da</span> <span class="o">*</span> <span class="n">epsilon</span><span class="p">,</span> <span class="n">dqkm1_da</span><span class="p">)</span>

        <span class="n">conditional</span> <span class="o">=</span> <span class="n">logicaland</span><span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">greater</span><span class="p">(</span><span class="n">grad_conditional</span><span class="p">,</span> <span class="n">epsilon</span><span class="p">))</span>

        <span class="k">return</span> <span class="p">(</span><span class="n">conditional</span><span class="p">,</span> <span class="n">select</span><span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">ans</span><span class="p">,</span> <span class="n">vals</span><span class="p">[</span><span class="mi">1</span><span class="p">]),</span> <span class="n">select</span><span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">t</span><span class="p">,</span> <span class="n">vals</span><span class="p">[</span><span class="mi">2</span><span class="p">]),</span>
                <span class="n">select</span><span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">y</span><span class="p">,</span> <span class="n">vals</span><span class="p">[</span><span class="mi">3</span><span class="p">]),</span> <span class="n">select</span><span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">z</span><span class="p">,</span> <span class="n">vals</span><span class="p">[</span><span class="mi">4</span><span class="p">]),</span>
                <span class="n">c</span><span class="p">,</span> <span class="n">select</span><span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">pkm1</span><span class="p">,</span> <span class="n">vals</span><span class="p">[</span><span class="mi">6</span><span class="p">]),</span>
                <span class="n">select</span><span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">qkm1</span><span class="p">,</span> <span class="n">vals</span><span class="p">[</span><span class="mi">7</span><span class="p">]),</span> <span class="n">select</span><span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">pkm2</span><span class="p">,</span> <span class="n">vals</span><span class="p">[</span><span class="mi">8</span><span class="p">]),</span>
                <span class="n">select</span><span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">qkm2</span><span class="p">,</span> <span class="n">vals</span><span class="p">[</span><span class="mi">9</span><span class="p">]),</span> <span class="n">select</span><span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">dpkm2_da</span><span class="p">,</span> <span class="n">vals</span><span class="p">[</span><span class="mi">10</span><span class="p">]),</span>
                <span class="n">select</span><span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">dqkm2_da</span><span class="p">,</span> <span class="n">vals</span><span class="p">[</span><span class="mi">11</span><span class="p">]),</span> <span class="n">select</span><span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">dpkm1_da</span><span class="p">,</span> <span class="n">vals</span><span class="p">[</span><span class="mi">12</span><span class="p">]),</span>
                <span class="n">select</span><span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">dqkm1_da</span><span class="p">,</span> <span class="n">vals</span><span class="p">[</span><span class="mi">13</span><span class="p">]),</span> <span class="n">select</span><span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">dans_da_new</span><span class="p">,</span> <span class="n">vals</span><span class="p">[</span><span class="mi">14</span><span class="p">]))</span>

    <span class="n">y</span> <span class="o">=</span> <span class="mi">1</span> <span class="o">-</span> <span class="n">a</span>
    <span class="n">z</span> <span class="o">=</span> <span class="n">x</span> <span class="o">+</span> <span class="n">y</span> <span class="o">+</span> <span class="mi">1</span>
    <span class="n">c</span> <span class="o">=</span> <span class="n">fill</span><span class="p">(</span><span class="n">dtype</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="n">shape</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="mi">0</span><span class="p">)</span>
    <span class="n">pkm2</span> <span class="o">=</span> <span class="n">fill</span><span class="p">(</span><span class="n">dtype</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="n">shape</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="mi">1</span><span class="p">)</span>
    <span class="n">qkm2</span> <span class="o">=</span> <span class="n">x</span>
    <span class="n">pkm1</span> <span class="o">=</span> <span class="n">x</span> <span class="o">+</span> <span class="mi">1</span>
    <span class="n">qkm1</span> <span class="o">=</span> <span class="n">z</span> <span class="o">*</span> <span class="n">x</span>
    <span class="n">ans</span> <span class="o">=</span> <span class="n">pkm1</span> <span class="o">/</span> <span class="n">qkm1</span>
    <span class="n">t</span> <span class="o">=</span> <span class="n">fill</span><span class="p">(</span><span class="n">dtype</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="n">shape</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="mi">1</span><span class="p">)</span>
    <span class="n">dpkm2_da</span> <span class="o">=</span> <span class="n">fill</span><span class="p">(</span><span class="n">dtype</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="n">shape</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="mi">0</span><span class="p">)</span>
    <span class="n">dqkm2_da</span> <span class="o">=</span> <span class="n">fill</span><span class="p">(</span><span class="n">dtype</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="n">shape</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="mi">0</span><span class="p">)</span>
    <span class="n">dpkm1_da</span> <span class="o">=</span> <span class="n">fill</span><span class="p">(</span><span class="n">dtype</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="n">shape</span><span class="p">(</span><span class="n">x</span><span class="p">),</span> <span class="mi">0</span><span class="p">)</span>
    <span class="n">dqkm1_da</span> <span class="o">=</span> <span class="o">-</span><span class="n">x</span>
    <span class="n">dans_da</span> <span class="o">=</span> <span class="p">(</span><span class="n">dpkm1_da</span> <span class="o">-</span> <span class="n">ans</span> <span class="o">*</span> <span class="n">dqkm1_da</span><span class="p">)</span> <span class="o">/</span> <span class="n">qkm1</span>
    <span class="n">vals</span> <span class="o">=</span> <span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">ans</span><span class="p">,</span> <span class="n">t</span><span class="p">,</span> <span class="n">y</span><span class="p">,</span> <span class="n">z</span><span class="p">,</span> <span class="n">c</span><span class="p">,</span> <span class="n">pkm1</span><span class="p">,</span> <span class="n">qkm1</span><span class="p">,</span> <span class="n">pkm2</span><span class="p">,</span> <span class="n">qkm2</span><span class="p">,</span> <span class="n">dpkm2_da</span><span class="p">,</span> <span class="n">dqkm2_da</span><span class="p">,</span> <span class="n">dpkm1_da</span><span class="p">,</span> <span class="n">dqkm1_da</span><span class="p">,</span> <span class="n">dans_da</span><span class="p">)</span>
    <span class="n">vals</span> <span class="o">=</span> <span class="n">_while_helper_func</span><span class="p">(</span><span class="n">cond</span><span class="p">,</span> <span class="n">body</span><span class="p">,</span> <span class="n">vals</span><span class="p">)</span>
    <span class="n">ans</span> <span class="o">=</span> <span class="n">vals</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
    <span class="k">return</span> <span class="n">ans</span> <span class="o">*</span> <span class="n">ax</span>


<span class="k">class</span> <span class="nc">IGamma</span><span class="p">(</span><span class="n">Cell</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Calculates lower regularized incomplete Gamma function.</span>
<span class="sd">    The lower regularized incomplete Gamma function is defined as:</span>

<span class="sd">    .. math::</span>
<span class="sd">        P(a, x) = gamma(a, x) / Gamma(a) = 1 - Q(a, x)</span>

<span class="sd">    where</span>

<span class="sd">    .. math::</span>
<span class="sd">        gamma(a, x) = \int_0^x t^{a-1} \exp^{-t} dt</span>

<span class="sd">    is the lower incomplete Gamma function.</span>

<span class="sd">    Above :math:`Q(a, x)` is the upper regularized complete Gamma function.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **a** (Tensor) - The input tensor. With float32 data type. `a` should have</span>
<span class="sd">          the same dtype with `x`.</span>
<span class="sd">        - **x** (Tensor) - The input tensor. With float32 data type. `x` should have</span>
<span class="sd">          the same dtype with `a`.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Tensor, has the same dtype as `a` and `x`.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If dtype of input x and a is not float16 nor float32,</span>
<span class="sd">                   or if x has different dtype with a.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; a = Tensor(np.array([2.0, 4.0, 6.0, 8.0]).astype(np.float32))</span>
<span class="sd">        &gt;&gt;&gt; x = Tensor(np.array([2.0, 3.0, 4.0, 5.0]).astype(np.float32))</span>
<span class="sd">        &gt;&gt;&gt; igamma = nn.IGamma()</span>
<span class="sd">        &gt;&gt;&gt; output = igamma(a, x)</span>
<span class="sd">        &gt;&gt;&gt; print (output)</span>
<span class="sd">        [0.593994  0.35276785  0.21486944  0.13337152]</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize IGamma.&quot;&quot;&quot;</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">IGamma</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
        <span class="c1"># const numbers</span>
        <span class="c1"># If more data types are supported, this float max value need to be selected.</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">log_maxfloat32</span> <span class="o">=</span> <span class="n">Tensor</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">finfo</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span><span class="o">.</span><span class="n">max</span><span class="p">),</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span>

        <span class="c1"># operations</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">logicaland</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">LogicalAnd</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">logicalor</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">LogicalOr</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">logicalnot</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">LogicalNot</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">equal</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Equal</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">greater</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Greater</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">less</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Less</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">neg</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Neg</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">log</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Log</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">exp</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Exp</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">select</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Select</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">zeroslike</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">ZerosLike</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">fill</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Fill</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">shape</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Shape</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">DType</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lgamma</span> <span class="o">=</span> <span class="n">LGamma</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">const</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">ScalarToArray</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">cast</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Cast</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">construct</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">a</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span>
        <span class="n">a_dtype</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">(</span><span class="n">a</span><span class="p">)</span>
        <span class="n">x_dtype</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
        <span class="n">_check_input_dtype</span><span class="p">(</span><span class="s2">&quot;a&quot;</span><span class="p">,</span> <span class="n">a_dtype</span><span class="p">,</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="n">_check_input_dtype</span><span class="p">(</span><span class="s2">&quot;x&quot;</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">,</span> <span class="n">a_dtype</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="n">domain_error</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">logicalor</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">less</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="mi">0</span><span class="p">),</span> <span class="bp">self</span><span class="o">.</span><span class="n">less</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="mi">0</span><span class="p">))</span>
        <span class="n">use_igammac</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">logicaland</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">greater</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="mi">1</span><span class="p">),</span> <span class="bp">self</span><span class="o">.</span><span class="n">greater</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">a</span><span class="p">))</span>
        <span class="n">ax</span> <span class="o">=</span> <span class="n">a</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">x</span><span class="p">)</span> <span class="o">-</span> <span class="n">x</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">lgamma</span><span class="p">(</span><span class="n">a</span><span class="p">)</span>
        <span class="n">para_shape</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">ax</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">para_shape</span> <span class="o">!=</span> <span class="p">():</span>
            <span class="n">broadcastto</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">BroadcastTo</span><span class="p">(</span><span class="n">para_shape</span><span class="p">)</span>
            <span class="n">x</span> <span class="o">=</span> <span class="n">broadcastto</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
            <span class="n">a</span> <span class="o">=</span> <span class="n">broadcastto</span><span class="p">(</span><span class="n">a</span><span class="p">)</span>
        <span class="n">x_is_zero</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">equal</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
        <span class="n">log_maxfloat</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">log_maxfloat32</span>
        <span class="n">underflow</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">less</span><span class="p">(</span><span class="n">ax</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">neg</span><span class="p">(</span><span class="n">log_maxfloat</span><span class="p">))</span>
        <span class="n">ax</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="n">ax</span><span class="p">)</span>
        <span class="n">enabled</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">logicalnot</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">logicalor</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">logicalor</span><span class="p">(</span><span class="n">x_is_zero</span><span class="p">,</span> <span class="n">domain_error</span><span class="p">),</span> <span class="n">underflow</span><span class="p">))</span>
        <span class="n">output</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">select</span><span class="p">(</span><span class="n">use_igammac</span><span class="p">,</span>
                             <span class="mi">1</span> <span class="o">-</span> <span class="n">_igammac_continued_fraction</span><span class="p">(</span><span class="n">ax</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">a</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">logicaland</span><span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="n">use_igammac</span><span class="p">)),</span>
                             <span class="n">_igamma_series</span><span class="p">(</span><span class="n">ax</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">a</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">logicaland</span><span class="p">(</span><span class="n">enabled</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">logicalnot</span><span class="p">(</span><span class="n">use_igammac</span><span class="p">))))</span>
        <span class="n">output</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">select</span><span class="p">(</span><span class="n">x_is_zero</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">zeroslike</span><span class="p">(</span><span class="n">output</span><span class="p">),</span> <span class="n">output</span><span class="p">)</span>
        <span class="n">output</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">select</span><span class="p">(</span><span class="n">domain_error</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">fill</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">(</span><span class="n">a</span><span class="p">),</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">a</span><span class="p">),</span> <span class="n">np</span><span class="o">.</span><span class="n">nan</span><span class="p">),</span> <span class="n">output</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">output</span>


<span class="k">class</span> <span class="nc">LBeta</span><span class="p">(</span><span class="n">Cell</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    This method avoids the numeric cancellation by explicitly</span>
<span class="sd">    decomposing lgamma into the Stirling approximation and an explicit log_gamma_correction, and cancelling</span>
<span class="sd">    the large terms from the Striling analytically.</span>

<span class="sd">    This is semantically equal to</span>

<span class="sd">    .. math::</span>
<span class="sd">        P(x, y) = lgamma(x) + lgamma(y) - lgamma(x + y).</span>

<span class="sd">    The method is more accurate for arguments above 8. The reason for accuracy loss in the naive computation</span>
<span class="sd">    is catastrophic cancellation between the lgammas.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **x** (Tensor) - The input tensor. With float16 or float32 data type. `x` should have</span>
<span class="sd">          the same dtype with `y`.</span>
<span class="sd">        - **y** (Tensor) - The input tensor. With float16 or float32 data type. `y` should have</span>
<span class="sd">          the same dtype with `x`.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Tensor, has the same dtype as `x` and `y`.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If dtype of `x` or `y` is neither float16 nor float32,</span>
<span class="sd">                   or if `x` has different dtype with `y`.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; x = Tensor(np.array([2.0, 4.0, 6.0, 8.0]).astype(np.float32))</span>
<span class="sd">        &gt;&gt;&gt; y = Tensor(np.array([2.0, 3.0, 14.0, 15.0]).astype(np.float32))</span>
<span class="sd">        &gt;&gt;&gt; lbeta = nn.LBeta()</span>
<span class="sd">        &gt;&gt;&gt; output = lbeta(y, x)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [-1.7917596  -4.094345  -12.000229  -14.754799]</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize LBeta.&quot;&quot;&quot;</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">LBeta</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
        <span class="c1"># const numbers</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">log_2pi</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="mi">2</span> <span class="o">*</span> <span class="n">np</span><span class="o">.</span><span class="n">pi</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">minimax_coeff</span> <span class="o">=</span> <span class="p">[</span><span class="o">-</span><span class="mf">0.165322962780713e-02</span><span class="p">,</span>
                              <span class="mf">0.837308034031215e-03</span><span class="p">,</span>
                              <span class="o">-</span><span class="mf">0.595202931351870e-03</span><span class="p">,</span>
                              <span class="mf">0.793650666825390e-03</span><span class="p">,</span>
                              <span class="o">-</span><span class="mf">0.277777777760991e-02</span><span class="p">,</span>
                              <span class="mf">0.833333333333333e-01</span><span class="p">]</span>

        <span class="c1"># operations</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">log</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Log</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">log1p</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Log1p</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">less</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Less</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">select</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Select</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">shape</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Shape</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">DType</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">lgamma</span> <span class="o">=</span> <span class="n">LGamma</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">const</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">ScalarToTensor</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">construct</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">y</span><span class="p">):</span>
        <span class="n">x_dtype</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
        <span class="n">y_dtype</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">(</span><span class="n">y</span><span class="p">)</span>
        <span class="n">_check_input_dtype</span><span class="p">(</span><span class="s2">&quot;x&quot;</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">,</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">float16</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="n">_check_input_dtype</span><span class="p">(</span><span class="s2">&quot;y&quot;</span><span class="p">,</span> <span class="n">y_dtype</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="n">x_plus_y</span> <span class="o">=</span> <span class="n">x</span> <span class="o">+</span> <span class="n">y</span>
        <span class="n">para_shape</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape</span><span class="p">(</span><span class="n">x_plus_y</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">para_shape</span> <span class="o">!=</span> <span class="p">():</span>
            <span class="n">broadcastto</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">BroadcastTo</span><span class="p">(</span><span class="n">para_shape</span><span class="p">)</span>
            <span class="n">x</span> <span class="o">=</span> <span class="n">broadcastto</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
            <span class="n">y</span> <span class="o">=</span> <span class="n">broadcastto</span><span class="p">(</span><span class="n">y</span><span class="p">)</span>
        <span class="n">comp_less</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">less</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">y</span><span class="p">)</span>
        <span class="n">x_min</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">select</span><span class="p">(</span><span class="n">comp_less</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">y</span><span class="p">)</span>
        <span class="n">y_max</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">select</span><span class="p">(</span><span class="n">comp_less</span><span class="p">,</span> <span class="n">y</span><span class="p">,</span> <span class="n">x</span><span class="p">)</span>

        <span class="nd">@constexpr</span>
        <span class="k">def</span> <span class="nf">_log_gamma_correction</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">minimax_coeff</span><span class="p">):</span>
            <span class="n">inverse_x</span> <span class="o">=</span> <span class="mf">1.</span> <span class="o">/</span> <span class="n">x</span>
            <span class="n">inverse_x_squared</span> <span class="o">=</span> <span class="n">inverse_x</span> <span class="o">*</span> <span class="n">inverse_x</span>
            <span class="n">accum</span> <span class="o">=</span> <span class="n">minimax_coeff</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
            <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">6</span><span class="p">):</span>
                <span class="n">accum</span> <span class="o">=</span> <span class="n">accum</span> <span class="o">*</span> <span class="n">inverse_x_squared</span> <span class="o">+</span> <span class="n">minimax_coeff</span><span class="p">[</span><span class="n">i</span><span class="p">]</span>
            <span class="k">return</span> <span class="n">accum</span> <span class="o">*</span> <span class="n">inverse_x</span>

        <span class="n">log_gamma_correction_x</span> <span class="o">=</span> <span class="n">_log_gamma_correction</span><span class="p">(</span><span class="n">x_min</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">minimax_coeff</span><span class="p">)</span>
        <span class="n">log_gamma_correction_y</span> <span class="o">=</span> <span class="n">_log_gamma_correction</span><span class="p">(</span><span class="n">y_max</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">minimax_coeff</span><span class="p">)</span>
        <span class="n">log_gamma_correction_x_y</span> <span class="o">=</span> <span class="n">_log_gamma_correction</span><span class="p">(</span><span class="n">x_plus_y</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">minimax_coeff</span><span class="p">)</span>

        <span class="c1"># Two large arguments case: y &gt;= x &gt;= 8.</span>
        <span class="n">log_beta_two_large</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">const</span><span class="p">(</span><span class="mf">0.5</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">log_2pi</span><span class="p">,</span> <span class="n">x_dtype</span><span class="p">)</span> <span class="o">-</span> <span class="mf">0.5</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">y_max</span><span class="p">)</span> \
                             <span class="o">+</span> <span class="n">log_gamma_correction_x</span> <span class="o">+</span> <span class="n">log_gamma_correction_y</span> <span class="o">-</span> <span class="n">log_gamma_correction_x_y</span> \
                             <span class="o">+</span> <span class="p">(</span><span class="n">x_min</span> <span class="o">-</span> <span class="mf">0.5</span><span class="p">)</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">x_min</span> <span class="o">/</span> <span class="p">(</span><span class="n">x_min</span> <span class="o">+</span> <span class="n">y_max</span><span class="p">))</span> <span class="o">-</span> <span class="n">y_max</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">log1p</span><span class="p">(</span><span class="n">x_min</span> <span class="o">/</span> <span class="n">y_max</span><span class="p">)</span>

        <span class="n">cancelled_stirling</span> <span class="o">=</span> <span class="o">-</span><span class="mi">1</span> <span class="o">*</span> <span class="p">(</span><span class="n">x_min</span> <span class="o">+</span> <span class="n">y_max</span> <span class="o">-</span> <span class="mf">0.5</span><span class="p">)</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">log1p</span><span class="p">(</span><span class="n">x_min</span> <span class="o">/</span> <span class="n">y_max</span><span class="p">)</span> <span class="o">-</span> <span class="n">x_min</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">y_max</span><span class="p">)</span> <span class="o">+</span> <span class="n">x_min</span>
        <span class="n">correction</span> <span class="o">=</span> <span class="n">log_gamma_correction_y</span> <span class="o">-</span> <span class="n">log_gamma_correction_x_y</span>
        <span class="n">log_gamma_difference_big_y</span> <span class="o">=</span> <span class="n">correction</span> <span class="o">+</span> <span class="n">cancelled_stirling</span>

        <span class="c1"># One large argument case: x &lt; 8, y &gt;= 8.</span>
        <span class="n">log_beta_one_large</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">lgamma</span><span class="p">(</span><span class="n">x_min</span><span class="p">)</span> <span class="o">+</span> <span class="n">log_gamma_difference_big_y</span>

        <span class="c1"># Small arguments case: x &lt;= y &lt; 8.</span>
        <span class="n">log_beta_small</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">lgamma</span><span class="p">(</span><span class="n">x_min</span><span class="p">)</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">lgamma</span><span class="p">(</span><span class="n">y_max</span><span class="p">)</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">lgamma</span><span class="p">(</span><span class="n">x_min</span> <span class="o">+</span> <span class="n">y_max</span><span class="p">)</span>
        <span class="n">comp_xless8</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">less</span><span class="p">(</span><span class="n">x_min</span><span class="p">,</span> <span class="mi">8</span><span class="p">)</span>
        <span class="n">comp_yless8</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">less</span><span class="p">(</span><span class="n">y_max</span><span class="p">,</span> <span class="mi">8</span><span class="p">)</span>
        <span class="n">temp</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">select</span><span class="p">(</span><span class="n">comp_yless8</span><span class="p">,</span> <span class="n">log_beta_small</span><span class="p">,</span> <span class="n">log_beta_one_large</span><span class="p">)</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">select</span><span class="p">(</span><span class="n">comp_xless8</span><span class="p">,</span> <span class="n">temp</span><span class="p">,</span> <span class="n">log_beta_two_large</span><span class="p">)</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">get_broadcast_matmul_shape</span><span class="p">(</span><span class="n">x_shape</span><span class="p">,</span> <span class="n">y_shape</span><span class="p">,</span> <span class="n">prim_name</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;get broadcast_matmul shape&quot;&quot;&quot;</span>
    <span class="n">msg_prefix</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="n">prim_name</span><span class="si">}</span><span class="s2">&#39;, the&quot;</span> <span class="k">if</span> <span class="n">prim_name</span> <span class="k">else</span> <span class="s2">&quot;The&quot;</span>
    <span class="k">if</span> <span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">x_shape</span><span class="p">)</span> <span class="o">&lt;</span> <span class="mi">2</span><span class="p">)</span> <span class="ow">or</span> <span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">y_shape</span><span class="p">)</span> <span class="o">&lt;</span> <span class="mi">2</span><span class="p">):</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> length of &#39;x_shape&#39; and &#39;y_shape&#39; should be equal to or greater than 2, &quot;</span>
                         <span class="sa">f</span><span class="s2">&quot;but got the length of &#39;x_shape&#39;: </span><span class="si">{</span><span class="nb">len</span><span class="p">(</span><span class="n">x_shape</span><span class="p">)</span><span class="si">}</span><span class="s2"> and the length of &#39;y_shape&#39;: &quot;</span>
                         <span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="nb">len</span><span class="p">(</span><span class="n">y_shape</span><span class="p">)</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>
    <span class="n">x_shape_batch</span> <span class="o">=</span> <span class="n">x_shape</span><span class="p">[:</span><span class="o">-</span><span class="mi">2</span><span class="p">]</span>
    <span class="n">y_shape_batch</span> <span class="o">=</span> <span class="n">y_shape</span><span class="p">[:</span><span class="o">-</span><span class="mi">2</span><span class="p">]</span>
    <span class="k">if</span> <span class="n">x_shape_batch</span> <span class="o">==</span> <span class="n">y_shape_batch</span><span class="p">:</span>
        <span class="k">return</span> <span class="n">x_shape</span><span class="p">,</span> <span class="n">y_shape</span>
    <span class="n">x_len</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">x_shape</span><span class="p">)</span>
    <span class="n">y_len</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">y_shape</span><span class="p">)</span>
    <span class="n">length</span> <span class="o">=</span> <span class="n">x_len</span> <span class="k">if</span> <span class="n">x_len</span> <span class="o">&lt;</span> <span class="n">y_len</span> <span class="k">else</span> <span class="n">y_len</span>
    <span class="n">broadcast_shape_back</span> <span class="o">=</span> <span class="p">[]</span>
    <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="o">-</span><span class="n">length</span><span class="p">,</span> <span class="o">-</span><span class="mi">2</span><span class="p">):</span>
        <span class="k">if</span> <span class="n">x_shape</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
            <span class="n">broadcast_shape_back</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">y_shape</span><span class="p">[</span><span class="n">i</span><span class="p">])</span>
        <span class="k">elif</span> <span class="n">y_shape</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
            <span class="n">broadcast_shape_back</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">x_shape</span><span class="p">[</span><span class="n">i</span><span class="p">])</span>
        <span class="k">elif</span> <span class="n">x_shape</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">==</span> <span class="n">y_shape</span><span class="p">[</span><span class="n">i</span><span class="p">]:</span>
            <span class="n">broadcast_shape_back</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">x_shape</span><span class="p">[</span><span class="n">i</span><span class="p">])</span>
        <span class="k">else</span><span class="p">:</span>
            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> &#39;x_shape[</span><span class="si">{</span><span class="n">i</span><span class="si">}</span><span class="s2">]&#39; should be equal to 1, or the &#39;y_shape[</span><span class="si">{</span><span class="n">i</span><span class="si">}</span><span class="s2">]&#39; should be equal &quot;</span>
                             <span class="sa">f</span><span class="s2">&quot;to 1, or the &#39;x_shape[</span><span class="si">{</span><span class="n">i</span><span class="si">}</span><span class="s2">]&#39; should be equal to &#39;y_shape[</span><span class="si">{</span><span class="n">i</span><span class="si">}</span><span class="s2">]&#39;, but got &quot;</span>
                             <span class="sa">f</span><span class="s2">&quot;&#39;x_shape[</span><span class="si">{</span><span class="n">i</span><span class="si">}</span><span class="s2">]&#39;: </span><span class="si">{</span><span class="n">x_shape</span><span class="p">[</span><span class="n">i</span><span class="p">]</span><span class="si">}</span><span class="s2">, &#39;y_shape[</span><span class="si">{</span><span class="n">i</span><span class="si">}</span><span class="s2">]&#39;: </span><span class="si">{</span><span class="n">y_shape</span><span class="p">[</span><span class="n">i</span><span class="p">]</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>

    <span class="n">broadcast_shape_front</span> <span class="o">=</span> <span class="n">y_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">:</span> <span class="n">y_len</span> <span class="o">-</span> <span class="n">length</span><span class="p">]</span> <span class="k">if</span> <span class="n">length</span> <span class="o">==</span> <span class="n">x_len</span> <span class="k">else</span> <span class="n">x_shape</span><span class="p">[</span><span class="mi">0</span><span class="p">:</span> <span class="n">x_len</span> <span class="o">-</span> <span class="n">length</span><span class="p">]</span>
    <span class="n">x_broadcast_shape</span> <span class="o">=</span> <span class="n">broadcast_shape_front</span> <span class="o">+</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">broadcast_shape_back</span><span class="p">)</span> <span class="o">+</span> <span class="n">x_shape</span><span class="p">[</span><span class="o">-</span><span class="mi">2</span><span class="p">:]</span>
    <span class="n">y_broadcast_shape</span> <span class="o">=</span> <span class="n">broadcast_shape_front</span> <span class="o">+</span> <span class="nb">tuple</span><span class="p">(</span><span class="n">broadcast_shape_back</span><span class="p">)</span> <span class="o">+</span> <span class="n">y_shape</span><span class="p">[</span><span class="o">-</span><span class="mi">2</span><span class="p">:]</span>
    <span class="k">return</span> <span class="n">x_broadcast_shape</span><span class="p">,</span> <span class="n">y_broadcast_shape</span>


<span class="nd">@constexpr</span>
<span class="k">def</span> <span class="nf">check_col_row_equal</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">,</span> <span class="n">x2_shape</span><span class="p">,</span> <span class="n">transpose_x1</span><span class="p">,</span> <span class="n">transpose_x2</span><span class="p">,</span> <span class="n">prim_name</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;check col and row equal&quot;&quot;&quot;</span>
    <span class="n">msg_prefix</span> <span class="o">=</span> <span class="sa">f</span><span class="s2">&quot;For &#39;</span><span class="si">{</span><span class="n">prim_name</span><span class="si">}</span><span class="s2">&#39;, the&quot;</span> <span class="k">if</span> <span class="n">prim_name</span> <span class="k">else</span> <span class="s2">&quot;The&quot;</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
        <span class="n">transpose_x1</span> <span class="o">=</span> <span class="kc">False</span>
        <span class="n">x1_shape</span> <span class="o">=</span> <span class="p">(</span><span class="mi">1</span><span class="p">,)</span> <span class="o">+</span> <span class="n">x1_shape</span>
    <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">x2_shape</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
        <span class="n">transpose_x2</span> <span class="o">=</span> <span class="kc">False</span>
        <span class="n">x2_shape</span> <span class="o">=</span> <span class="n">x2_shape</span> <span class="o">+</span> <span class="p">(</span><span class="mi">1</span><span class="p">,)</span>
    <span class="n">x1_last</span> <span class="o">=</span> <span class="n">x1_shape</span><span class="p">[</span><span class="o">-</span><span class="mi">2</span><span class="p">:]</span>
    <span class="n">x2_last</span> <span class="o">=</span> <span class="n">x2_shape</span><span class="p">[</span><span class="o">-</span><span class="mi">2</span><span class="p">:]</span>
    <span class="n">x1_col</span> <span class="o">=</span> <span class="n">x1_last</span><span class="p">[</span><span class="ow">not</span> <span class="n">transpose_x1</span><span class="p">]</span>  <span class="c1"># x1_col = x1_last[1] if (not transpose_a) else x1_last[0]</span>
    <span class="n">x2_row</span> <span class="o">=</span> <span class="n">x2_last</span><span class="p">[</span><span class="n">transpose_x2</span><span class="p">]</span>  <span class="c1"># x2_row = x2_last[0] if (not transpose_b) else x2_last[1]</span>
    <span class="k">if</span> <span class="n">x1_col</span> <span class="o">!=</span> <span class="n">x2_row</span><span class="p">:</span>
        <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">msg_prefix</span><span class="si">}</span><span class="s2"> column of matrix dimensions of &#39;x1&#39; should be equal to &quot;</span>
                         <span class="sa">f</span><span class="s2">&quot;the row of matrix dimensions of &#39;x2&#39;, but got &#39;x1_col&#39; </span><span class="si">{</span><span class="n">x1_col</span><span class="si">}</span><span class="s2"> and &#39;x2_row&#39; </span><span class="si">{</span><span class="n">x2_row</span><span class="si">}</span><span class="s2">.&quot;</span><span class="p">)</span>


<span class="k">def</span> <span class="nf">matmul_op_select</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">,</span> <span class="n">x2_shape</span><span class="p">,</span> <span class="n">transpose_x1</span><span class="p">,</span> <span class="n">transpose_x2</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;select matmul op&quot;&quot;&quot;</span>
    <span class="n">x1_dim</span><span class="p">,</span> <span class="n">x2_dim</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">),</span> <span class="nb">len</span><span class="p">(</span><span class="n">x2_shape</span><span class="p">)</span>
    <span class="k">if</span> <span class="n">x1_dim</span> <span class="o">==</span> <span class="mi">1</span> <span class="ow">and</span> <span class="n">x2_dim</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
        <span class="n">matmul_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Mul</span><span class="p">()</span>
    <span class="k">elif</span> <span class="n">x1_dim</span> <span class="o">&lt;=</span> <span class="mi">2</span> <span class="ow">and</span> <span class="n">x2_dim</span> <span class="o">&lt;=</span> <span class="mi">2</span><span class="p">:</span>
        <span class="n">transpose_x1</span> <span class="o">=</span> <span class="kc">False</span> <span class="k">if</span> <span class="n">x1_dim</span> <span class="o">==</span> <span class="mi">1</span> <span class="k">else</span> <span class="n">transpose_x1</span>
        <span class="n">transpose_x2</span> <span class="o">=</span> <span class="kc">False</span> <span class="k">if</span> <span class="n">x2_dim</span> <span class="o">==</span> <span class="mi">1</span> <span class="k">else</span> <span class="n">transpose_x2</span>
        <span class="n">matmul_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">MatMul</span><span class="p">(</span><span class="n">transpose_x1</span><span class="p">,</span> <span class="n">transpose_x2</span><span class="p">)</span>
    <span class="k">elif</span> <span class="n">x1_dim</span> <span class="o">==</span> <span class="mi">1</span> <span class="ow">and</span> <span class="n">x2_dim</span> <span class="o">&gt;</span> <span class="mi">2</span><span class="p">:</span>
        <span class="n">matmul_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">BatchMatMul</span><span class="p">(</span><span class="kc">False</span><span class="p">,</span> <span class="n">transpose_x2</span><span class="p">)</span>
    <span class="k">elif</span> <span class="n">x1_dim</span> <span class="o">&gt;</span> <span class="mi">2</span> <span class="ow">and</span> <span class="n">x2_dim</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
        <span class="n">matmul_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">BatchMatMul</span><span class="p">(</span><span class="n">transpose_x1</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>
    <span class="k">else</span><span class="p">:</span>
        <span class="n">matmul_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">BatchMatMul</span><span class="p">(</span><span class="n">transpose_x1</span><span class="p">,</span> <span class="n">transpose_x2</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">matmul_op</span>


<div class="viewcode-block" id="MatMul"><a class="viewcode-back" href="../../../../api_python/nn/mindspore.nn.MatMul.html#mindspore.nn.MatMul">[docs]</a><span class="k">class</span> <span class="nc">MatMul</span><span class="p">(</span><span class="n">Cell</span><span class="p">):</span>
    <span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    The nn.MatMul interface is deprecated, please use the :class:`mindspore.ops.matmul` instead.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        deprecated</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="nd">@deprecated</span><span class="p">(</span><span class="s1">&#39;1.2&#39;</span><span class="p">,</span> <span class="s1">&#39;ops.matmul&#39;</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">transpose_x1</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">transpose_x2</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize MatMul.&quot;&quot;&quot;</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">MatMul</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>

        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;transpose_x1&#39;</span><span class="p">,</span> <span class="n">transpose_x1</span><span class="p">,</span> <span class="p">[</span><span class="nb">bool</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;transpose_x2&#39;</span><span class="p">,</span> <span class="n">transpose_x2</span><span class="p">,</span> <span class="p">[</span><span class="nb">bool</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">transpose_x1</span> <span class="o">=</span> <span class="n">transpose_x1</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">transpose_x2</span> <span class="o">=</span> <span class="n">transpose_x2</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">shape_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Shape</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">expand_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">ExpandDims</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">squeeze_left_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Squeeze</span><span class="p">(</span><span class="o">-</span><span class="mi">2</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">squeeze_right_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Squeeze</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">reduce_sum_op</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">ReduceSum</span><span class="p">(</span><span class="n">keep_dims</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">construct</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x1</span><span class="p">,</span> <span class="n">x2</span><span class="p">):</span>
        <span class="n">x1_shape</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape_op</span><span class="p">(</span><span class="n">x1</span><span class="p">)</span>
        <span class="n">x2_shape</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape_op</span><span class="p">(</span><span class="n">x2</span><span class="p">)</span>
        <span class="n">check_col_row_equal</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">,</span> <span class="n">x2_shape</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">transpose_x1</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">transpose_x2</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="n">matmul_op</span> <span class="o">=</span> <span class="n">matmul_op_select</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">,</span> <span class="n">x2_shape</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">transpose_x1</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">transpose_x2</span><span class="p">)</span>

        <span class="n">x1_dim</span><span class="p">,</span> <span class="n">x2_dim</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">),</span> <span class="nb">len</span><span class="p">(</span><span class="n">x2_shape</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">x1_dim</span> <span class="o">==</span> <span class="n">x2_dim</span> <span class="ow">and</span> <span class="n">x2_dim</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
            <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">reduce_sum_op</span><span class="p">(</span><span class="n">matmul_op</span><span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="n">x2</span><span class="p">),</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">x1_dim</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
            <span class="n">x1</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">expand_op</span><span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="mi">0</span><span class="p">)</span>
            <span class="n">x1_shape</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape_op</span><span class="p">(</span><span class="n">x1</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">x2_dim</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
            <span class="n">x2</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">expand_op</span><span class="p">(</span><span class="n">x2</span><span class="p">,</span> <span class="mi">1</span><span class="p">)</span>
            <span class="n">x2_shape</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">shape_op</span><span class="p">(</span><span class="n">x2</span><span class="p">)</span>

        <span class="n">x1_broadcast_shape</span><span class="p">,</span> <span class="n">x2_broadcast_shape</span> <span class="o">=</span> <span class="n">get_broadcast_matmul_shape</span><span class="p">(</span><span class="n">x1_shape</span><span class="p">,</span> <span class="n">x2_shape</span><span class="p">)</span>
        <span class="n">x1_broadcast_to</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">BroadcastTo</span><span class="p">(</span><span class="n">x1_broadcast_shape</span><span class="p">)</span>
        <span class="n">x2_broadcast_to</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">BroadcastTo</span><span class="p">(</span><span class="n">x2_broadcast_shape</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">x1_broadcast_shape</span> <span class="o">!=</span> <span class="n">x1_shape</span><span class="p">:</span>
            <span class="n">x1</span> <span class="o">=</span> <span class="n">x1_broadcast_to</span><span class="p">(</span><span class="n">x1</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">x2_broadcast_shape</span> <span class="o">!=</span> <span class="n">x2_shape</span><span class="p">:</span>
            <span class="n">x2</span> <span class="o">=</span> <span class="n">x2_broadcast_to</span><span class="p">(</span><span class="n">x2</span><span class="p">)</span>

        <span class="n">matmul_broadcast</span> <span class="o">=</span> <span class="n">matmul_op</span><span class="p">(</span><span class="n">x1</span><span class="p">,</span> <span class="n">x2</span><span class="p">)</span>

        <span class="k">if</span> <span class="n">x1_dim</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
            <span class="n">matmul_broadcast</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">squeeze_left_op</span><span class="p">(</span><span class="n">matmul_broadcast</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">x2_dim</span> <span class="o">==</span> <span class="mi">1</span><span class="p">:</span>
            <span class="n">matmul_broadcast</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">squeeze_right_op</span><span class="p">(</span><span class="n">matmul_broadcast</span><span class="p">)</span>

        <span class="k">return</span> <span class="n">matmul_broadcast</span></div>


<span class="k">class</span> <span class="nc">Moments</span><span class="p">(</span><span class="n">Cell</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Calculates the mean and variance of `x`.</span>

<span class="sd">    The mean and variance are calculated by aggregating the contents of `input_x` across axes.</span>
<span class="sd">    If `input_x` is 1-D and axes = [0] this is just the mean and variance of a vector.</span>

<span class="sd">    Args:</span>
<span class="sd">        axis (Union[int, tuple(int)]): Calculates the mean and variance along the specified axis. Default: None.</span>
<span class="sd">        keep_dims (bool): If true, The dimension of mean and variance are identical with input&#39;s.</span>
<span class="sd">                          If false, don&#39;t keep these dimensions. Default: None.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **x** (Tensor) - The tensor to be calculated. Only float16 and float32 are supported.</span>
<span class="sd">          :math:`(N,*)` where :math:`*` means,any number of additional dimensions.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        - **mean** (Tensor) - The mean of `x`, with the same data type as input `x`.</span>
<span class="sd">        - **variance** (Tensor) - The variance of `x`, with the same data type as input `x`.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If `axis` is not one of int, tuple, None.</span>
<span class="sd">        TypeError: If `keep_dims` is neither bool nor None.</span>
<span class="sd">        TypeError: If dtype of `x` is neither float16 nor float32.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``Ascend`` ``GPU`` ``CPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; x = Tensor(np.array([[[[1, 2, 3, 4], [3, 4, 5, 6]]]]), mindspore.float32)</span>
<span class="sd">        &gt;&gt;&gt; net = nn.Moments(axis=0, keep_dims=True)</span>
<span class="sd">        &gt;&gt;&gt; output = net(x)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        (Tensor(shape=[1, 1, 2, 4], dtype=Float32, value=</span>
<span class="sd">        [[[[ 1.00000000e+00, 2.00000000e+00, 3.00000000e+00, 4.00000000e+00],</span>
<span class="sd">           [ 3.00000000e+00, 4.00000000e+00, 5.00000000e+00, 6.00000000e+00]]]]),</span>
<span class="sd">        Tensor(shape=[1, 1, 2, 4], dtype=Float32, value=</span>
<span class="sd">        [[[[ 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00],</span>
<span class="sd">           [ 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00]]]]))</span>
<span class="sd">        &gt;&gt;&gt; net = nn.Moments(axis=1, keep_dims=True)</span>
<span class="sd">        &gt;&gt;&gt; output = net(x)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        (Tensor(shape=[1, 1, 2, 4], dtype=Float32, value=</span>
<span class="sd">        [[[[ 1.00000000e+00, 2.00000000e+00, 3.00000000e+00, 4.00000000e+00],</span>
<span class="sd">           [ 3.00000000e+00, 4.00000000e+00, 5.00000000e+00, 6.00000000e+00]]]]),</span>
<span class="sd">        Tensor(shape=[1, 1, 2, 4], dtype=Float32, value=</span>
<span class="sd">        [[[[ 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00],</span>
<span class="sd">           [ 0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00]]]]))</span>
<span class="sd">        &gt;&gt;&gt; net = nn.Moments(axis=2, keep_dims=True)</span>
<span class="sd">        &gt;&gt;&gt; output = net(x)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        (Tensor(shape=[1, 1, 1, 4], dtype=Float32, value=</span>
<span class="sd">        [[[[ 2.00000000e+00, 3.00000000e+00, 4.00000000e+00, 5.00000000e+00]]]]),</span>
<span class="sd">        Tensor(shape=[1, 1, 1, 4], dtype=Float32, value=</span>
<span class="sd">        [[[[ 1.00000000e+00, 1.00000000e+00, 1.00000000e+00, 1.00000000e+00]]]]))</span>
<span class="sd">        &gt;&gt;&gt; net = nn.Moments(axis=3, keep_dims=True)</span>
<span class="sd">        &gt;&gt;&gt; output = net(x)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        (Tensor(shape=[1, 1, 2, 1], dtype=Float32, value=</span>
<span class="sd">        [[[[ 2.50000000e+00],</span>
<span class="sd">           [ 4.50000000e+00]]]]), Tensor(shape=[1, 1, 2, 1], dtype=Float32, value=</span>
<span class="sd">        [[[[ 1.25000000e+00],</span>
<span class="sd">           [ 1.25000000e+00]]]]))</span>
<span class="sd">    &quot;&quot;&quot;</span>

    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">keep_dims</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize Moments.&quot;&quot;&quot;</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">Moments</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
        <span class="k">if</span> <span class="n">axis</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">axis</span> <span class="o">=</span> <span class="p">()</span>
        <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">axis</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">):</span>
            <span class="k">for</span> <span class="n">idx</span><span class="p">,</span> <span class="n">item</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">axis</span><span class="p">):</span>
                <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s2">&quot;axis[</span><span class="si">%d</span><span class="s2">]&quot;</span> <span class="o">%</span> <span class="n">idx</span><span class="p">,</span> <span class="n">item</span><span class="p">,</span> <span class="p">[</span><span class="nb">int</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">axis</span> <span class="o">=</span> <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;axis&#39;</span><span class="p">,</span> <span class="n">axis</span><span class="p">,</span> <span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">tuple</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">keep_dims</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
            <span class="n">keep_dims</span> <span class="o">=</span> <span class="kc">False</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">keep_dims</span> <span class="o">=</span> <span class="n">validator</span><span class="o">.</span><span class="n">check_value_type</span><span class="p">(</span><span class="s1">&#39;keep_dims&#39;</span><span class="p">,</span> <span class="n">keep_dims</span><span class="p">,</span> <span class="p">[</span><span class="nb">bool</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">cast</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Cast</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">reduce_mean</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">ReduceMean</span><span class="p">(</span><span class="n">keep_dims</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">square_diff</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">SquaredDifference</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">squeeze</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Squeeze</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">axis</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">construct</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span>
        <span class="n">tensor_dtype</span> <span class="o">=</span> <span class="n">F</span><span class="o">.</span><span class="n">dtype</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
        <span class="n">_check_input_dtype</span><span class="p">(</span><span class="s2">&quot;input x&quot;</span><span class="p">,</span> <span class="n">tensor_dtype</span><span class="p">,</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">float16</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">tensor_dtype</span> <span class="o">==</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float16</span><span class="p">:</span>
            <span class="n">x</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">cast</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span>
        <span class="n">mean</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">reduce_mean</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">axis</span><span class="p">)</span>
        <span class="n">variance</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">reduce_mean</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">square_diff</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">F</span><span class="o">.</span><span class="n">stop_gradient</span><span class="p">(</span><span class="n">mean</span><span class="p">)),</span> <span class="bp">self</span><span class="o">.</span><span class="n">axis</span><span class="p">)</span>
        <span class="k">if</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">keep_dims</span><span class="p">:</span>
            <span class="n">mean</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">squeeze</span><span class="p">(</span><span class="n">mean</span><span class="p">)</span>
            <span class="n">variance</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">squeeze</span><span class="p">(</span><span class="n">variance</span><span class="p">)</span>
        <span class="k">if</span> <span class="n">tensor_dtype</span> <span class="o">==</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float16</span><span class="p">:</span>
            <span class="n">mean</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">cast</span><span class="p">(</span><span class="n">mean</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float16</span><span class="p">)</span>
            <span class="n">variance</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">cast</span><span class="p">(</span><span class="n">variance</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float16</span><span class="p">)</span>
            <span class="k">return</span> <span class="n">mean</span><span class="p">,</span> <span class="n">variance</span>
        <span class="k">return</span> <span class="n">mean</span><span class="p">,</span> <span class="n">variance</span>


<span class="k">class</span> <span class="nc">MatInverse</span><span class="p">(</span><span class="n">Cell</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Calculates the inverse of Positive-Definite Hermitian matrix using Cholesky decomposition.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **x** (Tensor[Number]) - The input tensor. It must be a positive-definite matrix.</span>
<span class="sd">          With float16 or float32 data type.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Tensor, has the same dtype as the `x`.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If dtype of `x` is neither float16 nor float32.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``GPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; x = Tensor(np.array([[4, 12, -16], [12, 37, -43], [-16, -43, 98]]).astype(np.float32))</span>
<span class="sd">        &gt;&gt;&gt; op = nn.MatInverse()</span>
<span class="sd">        &gt;&gt;&gt; output = op(x)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        [[49.36112  -13.555558  2.1111116]</span>
<span class="sd">         [-13.555558  3.7777784  -0.5555557]</span>
<span class="sd">         [2.1111116  -0.5555557  0.11111113]]</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize MatInverse.&quot;&quot;&quot;</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">MatInverse</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">DType</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">choleskytrsm</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">CholeskyTrsm</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">matmul</span> <span class="o">=</span> <span class="n">MatMul</span><span class="p">(</span><span class="n">transpose_x1</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>

    <span class="k">def</span> <span class="nf">construct</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">a</span><span class="p">):</span>
        <span class="n">input_dtype</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">(</span><span class="n">a</span><span class="p">)</span>
        <span class="n">_check_input_dtype</span><span class="p">(</span><span class="s2">&quot;input_a&quot;</span><span class="p">,</span> <span class="n">input_dtype</span><span class="p">,</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">float16</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="n">l_inverse</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">choleskytrsm</span><span class="p">(</span><span class="n">a</span><span class="p">)</span>
        <span class="n">a_inverse</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">matmul</span><span class="p">(</span><span class="n">l_inverse</span><span class="p">,</span> <span class="n">l_inverse</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">a_inverse</span>


<span class="k">class</span> <span class="nc">MatDet</span><span class="p">(</span><span class="n">Cell</span><span class="p">):</span>
    <span class="sd">&quot;&quot;&quot;</span>
<span class="sd">    Calculates the determinant of Positive-Definite Hermitian matrix using Cholesky decomposition.</span>

<span class="sd">    Inputs:</span>
<span class="sd">        - **x** (Tensor[Number]) - The input tensor. It must be a positive-definite matrix.</span>
<span class="sd">          With float16 or float32 data type.</span>

<span class="sd">    Outputs:</span>
<span class="sd">        Tensor, has the same dtype as the `x`.</span>

<span class="sd">    Raises:</span>
<span class="sd">        TypeError: If dtype of `x` is neither float16 nor float32.</span>

<span class="sd">    Supported Platforms:</span>
<span class="sd">        ``GPU``</span>

<span class="sd">    Examples:</span>
<span class="sd">        &gt;&gt;&gt; x = Tensor(np.array([[4, 12, -16], [12, 37, -43], [-16, -43, 98]]).astype(np.float32))</span>
<span class="sd">        &gt;&gt;&gt; op = nn.MatDet()</span>
<span class="sd">        &gt;&gt;&gt; output = op(x)</span>
<span class="sd">        &gt;&gt;&gt; print(output)</span>
<span class="sd">        35.999996</span>
<span class="sd">    &quot;&quot;&quot;</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
        <span class="sd">&quot;&quot;&quot;Initialize MatDet.&quot;&quot;&quot;</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">MatDet</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">DType</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">cholesky</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Cholesky</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">det_triangle</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">DetTriangle</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">square</span> <span class="o">=</span> <span class="n">P</span><span class="o">.</span><span class="n">Square</span><span class="p">()</span>

    <span class="k">def</span> <span class="nf">construct</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">a</span><span class="p">):</span>
        <span class="n">input_dtype</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">(</span><span class="n">a</span><span class="p">)</span>
        <span class="n">_check_input_dtype</span><span class="p">(</span><span class="s2">&quot;input_a&quot;</span><span class="p">,</span> <span class="n">input_dtype</span><span class="p">,</span> <span class="p">[</span><span class="n">mstype</span><span class="o">.</span><span class="n">float16</span><span class="p">,</span> <span class="n">mstype</span><span class="o">.</span><span class="n">float32</span><span class="p">],</span> <span class="bp">self</span><span class="o">.</span><span class="n">cls_name</span><span class="p">)</span>
        <span class="n">l</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">cholesky</span><span class="p">(</span><span class="n">a</span><span class="p">)</span>
        <span class="n">l_det</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">det_triangle</span><span class="p">(</span><span class="n">l</span><span class="p">)</span>
        <span class="n">a_det</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">square</span><span class="p">(</span><span class="n">l_det</span><span class="p">)</span>
        <span class="k">return</span> <span class="n">a_det</span>
</pre></div>

           </div>
           
          </div>
          <footer>

  <hr/>

  <div role="contentinfo">
    <p>
        &#169; Copyright 2021, MindSpore.

    </p>
  </div>
    
    
    
    Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
    
    <a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
    
    provided by <a href="https://readthedocs.org">Read the Docs</a>. 

</footer>
        </div>
      </div>

    </section>

  </div>
  

  <script type="text/javascript">
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script>

  
  
    
   

</body>
</html>